package com.ferry.lock.engine.zookeeper.menagerie;

import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooKeeper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

/**
 * Created by daiyong
 */
public class ZKSessionPoller {

	private static final Logger logger = LoggerFactory.getLogger(ZKSessionPoller.class);

	private final long pollIntervalMs;

	private final ZooKeeper zk;

	private final Object disconnectTimeLock = "Lock";

	private Long startDisconnectTime;

	private final ConnectionListener pollListener;

	private final ScheduledExecutorService poller = Executors.newScheduledThreadPool(1, new ThreadFactory() {
		@Override
		public Thread newThread(Runnable runnable) {
			Thread t = new Thread(runnable);
			t.setName("ZKSessionPoller");
			return t;
		}
	});

	public ZKSessionPoller(ZooKeeper zk, long pollIntervalMs, ConnectionListener pollListener) {
		this.pollIntervalMs = pollIntervalMs;
		this.zk = zk;
		this.pollListener = pollListener;
	}

	public void startPolling(){
		poller.scheduleWithFixedDelay(new SessionPoller(), 0l, pollIntervalMs, TimeUnit.MILLISECONDS);
	}

	public void stopPolling(){
		poller.shutdownNow();
	}

	private class SessionPoller implements Runnable{
		private final int sessionTimeoutPeriod;

		private SessionPoller() {
			sessionTimeoutPeriod = zk.getSessionTimeout();
		}

		@Override
		public void run() {
			if(Thread.currentThread().isInterrupted())return; //we've been canceled, so return
			if(logger.isTraceEnabled())
				logger.trace("current state of ZooKeeper object: "+ zk.getState());

			try{
				zk.exists("/", false);
				synchronized(disconnectTimeLock){
					startDisconnectTime = null;
				}
			} catch (InterruptedException e) {
				e.printStackTrace();
			} catch (KeeperException e) {
				if(e.code() == KeeperException.Code.SESSIONEXPIRED){
					expire();
				}else if(e.code() == KeeperException.Code.CONNECTIONLOSS){
					logger.debug("Received a ConnectionLoss Exception, determining if our session has expired");
					long currentTime = System.currentTimeMillis();
					boolean shouldExpire = false;
					synchronized(disconnectTimeLock){
						if(startDisconnectTime == null){
							startDisconnectTime = currentTime;
						} else if ((currentTime - startDisconnectTime) > sessionTimeoutPeriod){
							shouldExpire=true;
						}
					}
					if(shouldExpire)expire();
				}else{
					e.printStackTrace();
				}
			}
		}
	}

	private void expire() {
		//session expired!
		logger.info("Session has expired, notifying listenerand shutting down poller");
		ZKSessionPoller.this.stopPolling();
		pollListener.expired();
	}

}
